import torch
import torch.nn.functional as F
import torch.nn as nn

from .sa import SA
from .bconv import Bconv
from .resnet import ResNet18,init_weight
from .rf import RF
from .pdc import PDC

class SINet(nn.Module):
    def __init__(self,RF_out_channels=32):
        super(SINet, self).__init__()
        # 上采样和下采样
        self.downsample2 = nn.MaxPool2d(2, stride=2)
        self.upsample2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.upsample4 = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True)
        self.upsample8 = nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True)

        # resnent18
        self.resnet = ResNet18()   # 3->64->64 -> 128 -> 256 -> 512
        init_weight(self.resnet)

        # SM的RF和PDC
        self.rf4_sm = RF(128, RF_out_channels)
        self.rf3_sm = RF(896, RF_out_channels)
        self.rf2_sm = RF(768, RF_out_channels)
        self.rf1_sm = RF(512, RF_out_channels)
        self.pdc_sm = PDC(RF_out_channels, f4=True)

        # IM的RF和PDC
        self.sa = SA(15,4)
        self.rf1_im = RF(512, RF_out_channels)
        self.rf2_im = RF(256, RF_out_channels)
        self.rf3_im = RF(128, RF_out_channels)
        self.pdc_im = PDC(RF_out_channels, f4=False)

    def forward(self,x):
        resdict = self.resnet(x)
        x0 = resdict['X0']
        x1 = resdict['X1']
        x2 = resdict['X2']
        x3 = resdict['X3']
        x4 = resdict['X4']

        x01_cat = torch.cat((x0, x1), dim=1)
        x01_cat = self.downsample2(x01_cat)
        sm_rf4 = self.rf4_sm(x01_cat)

        x234_cat = torch.cat((x2,self.upsample2(x3),self.upsample4(x4)),dim=1)
        sm_rf3 = self.rf3_sm(x234_cat)

        x34_cat = torch.cat((x3,self.upsample2(x4)),dim=1)
        sm_rf2 = self.rf2_sm(x34_cat)
        sm_rf1 = self.rf1_sm(x4)

        out1 = self.pdc_sm(sm_rf1,sm_rf2,sm_rf3,sm_rf4)

        sa_out = self.sa(out1.sigmoid(),x2)

        x3_1 = self.resnet.layer3_im(sa_out)
        x4_1 = self.resnet.layer4_im(x3_1)
        
        im_rf1 = self.rf1_im(x4_1)
        im_rf2 = self.rf2_im(x3_1)
        im_rf3 = self.rf3_im(sa_out)
        out2 = self.pdc_im(im_rf1,im_rf2,im_rf3)

        out1 = self.upsample8(out1)
        out2 = self.upsample8(out2)

        return out1, out2
    
    
if __name__ == '__main__':
    model = SINet()
    x = torch.randn(1,3,352,352)
    out1, out2 = model(x)
    print(out1.shape)
    print(out2.shape)
